Polychromatic Propagations

Prysm has a long heritage solving the monochromatic problem very quickly. However, it used a brute force ‘propagate and interpolate’ approach to solving the polychromatic problem. v0.19 offers large speedup by using matrix triple product DFTs to perform polychromatic propagations. This results in forward model times that are significantly faster, and clearer code when propagating to a grid for a detector.

This notebook shows in the fewest possible lines the speedup by using a zero OPD circular pupil under nine discrete wavelengths. It also shows propagating to a detector grid.

from operator import add
from functools import reduce

import numpy as np

from matplotlib import pyplot as plt

from prysm import Pupil, PSF
from prysm.propagation import focus_units, Wavefront

wvls = np.linspace(.4, .7, 9) # 400 to 900 nm, 9 wavelengths
pups = [Pupil(wavelength=w, samples=512) for w in wvls]

The old way

%%timeit
psfs = [PSF.from_pupil(pup, efl=1, Q=2) for pup in pups]
poly_psf = PSF.polychromatic(psfs)

2.32 s ± 121 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# the new way - setup for equivalence of output
x, y = focus_units(pups[0].fcn, pups[0].sample_spacing, 1, wvls[0], 2)
sample_spacing = x[1] - x[0]

The new way

%%timeit
psf_fields = [p.astype(Wavefront)\
              .focus_fixed_sampling(efl=1, sample_spacing=sample_spacing, samples=1024) for p in pups]
psf_intensities = [abs(w.fcn)**2 for w in psf_fields]
poly_psf2 = PSF(x=x, y=y, data=reduce(add, psf_intensities))

458 ms ± 6.07 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In this simple example, the new way is about 4x faster. In this case, the simulation was done at Q=2 for all colors in the ‘old’ polychromatic way. This results in some numerical errors, where the new way is error free. At larger Qs the old way will have improved accuracy, but also increased computation time. To show the true power of the new way, we will compare old and new for Q=8, and use the flexibility of the matrix triple product to compute a smaller output domain:

the old (high oversampling):

%%timeit
psfs = [PSF.from_pupil(pup, efl=1, Q=8) for pup in pups]
poly_psf3 = PSF.polychromatic(psfs)

48.8 s ± 265 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# the new way - setup for equivalence of output
x, y = focus_units(pups[0].fcn, pups[0].sample_spacing, 1, wvls[0], 8)
sample_spacing = x[1] - x[0]

the new (high oversampling)

%%timeit
psf_fields = [p.astype(Wavefront)\
              .focus_fixed_sampling(efl=1, sample_spacing=sample_spacing, samples=256) for p in pups]
psf_intensities = [abs(w.fcn)**2 for w in psf_fields]
poly_psf4 = PSF(x=x, y=y, data=reduce(add, psf_intensities))

103 ms ± 4.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

The parameters of this second example are not particularly relevant outside coronagraphy or simulation study of astronomical instrumentation due to the large Q, but we see a nearly 500x speedup for use of the matrix triple product.

While the output data is not strictly the same since the matrix triple product is computed over a smaller domain, their value is the same since we do not care about the region far from the core of the PSF. This speedup allows computations that may require a supercomputer to be done on a laptop.